# -*- coding: utf-8 -*-  
'''
tbert的损失函数，交叉熵

Created on 2021年9月12日
@author: luoyi
'''
import tensorflow as tf


#    tbert的损失
class TBertLoss(tf.losses.Loss):
    '''交叉熵'''
    def __init__(self,
                 name='tbert-loss',
                 **kwargs):
        super(TBertLoss, self).__init__(name=name, **kwargs)
        pass
    
    def call(self, y_true, y_pred):
        '''
            @param y_true: Tensor(batch_size, )     真实主题id
            @param y_pred: Tensor(batch_size, K)    每个主题的概率
        '''
        #    批量大小
        B = tf.math.count_nonzero(y_true[:, 0] >= 0)
        
        #    每篇文章的真实主题的预测概率
        idx_B = tf.range(B, dtype=tf.int64)
        idx_t = tf.squeeze(y_true, axis=-1)
        idx = tf.stack([idx_B, idx_t], axis=-1)
        y_pred = tf.gather_nd(params=y_pred, indices=idx)       #    Tensor(batch_size, )
        
        #    cross entropy
        loss = -tf.math.log(y_pred)                             #    Tensor(batch_size, )
        
        return loss
    
    pass